##!/opt/local/bin/python

from rfvectors import cavity
import csv
import matplotlib.pyplot as plt
from numpy import nan, arange, array
from scipy.optimize import fmin, fminbound

def linacreflected(Ql, filename='spokenewshort.csv', f=704.42e6, delF=nan, I_b0=50e-3, I_b=50e-3, RoverQ=477.0, 
        beta_0=0.92, Eampfile=nan, verbose=False):
    "A function to return the total reflected power from a linac."
    if verbose: print "Ql = %e :: file = %s :: Eamp file = %s :: Beta_0 = %f :: RoverQ = %f" % (Ql, filename, Eampfile, beta_0, RoverQ)
    data = csv.reader(open(filename, 'rU'), dialect=csv.excel_tab, delimiter=',')
    totalRef = 0
    for row in data:
        V = float(row[0]) * float(row[3])
        ph = float(row[1])
        b = float(row[2])
        C = cavity(f=f, delF=delF, I_b0=I_b0, I_b=I_b, V_cav=V, beamph=ph, 
                RoverQ=RoverQ, Ql=Ql, beta_0=beta_0, beta=b, Eampfile=Eampfile)
        Pref = C.Pg - C.Pbeam
        totalRef += Pref
    return totalRef

if __name__=="__main__":
    # Choose a Ql for each SC linac section to minimise the total reflected power of that section
    Qlopt = {'spoke': 0, 'medbeta': 0, 'highbeta': 0}
    spokeargs    = ('spokenewshort.csv',    352.21e6, nan, 50e-3, 50e-3, 500.0, 0.46, 'Ezvszspoke.dat', False)
    medbetaargs  = ('lowbetanew.csv',  704.42e6, nan, 50e-3, 50e-3, 300.0, 0.70, nan, False)
    highbetaargs = ('highbetanew.csv', 704.42e6, nan, 50e-3, 50e-3, 477.0, 0.92, nan, False)
    Qlopt['spoke']    = fminbound(linacreflected, 1e4, 1e7, args=spokeargs)
    Qlopt['medbeta']  = fminbound(linacreflected, 1e4, 1e7, args=medbetaargs)
    Qlopt['highbeta'] = fminbound(linacreflected, 1e4, 1e7, args=highbetaargs)
    print "Optimized Q_l (spoke)     = %0.3e" % Qlopt['spoke']
    print "Optimized Q_l (med beta)  = %0.3e" % Qlopt['medbeta']
    print "Optimized Q_l (high beta) = %0.3e" % Qlopt['highbeta']
    
    cavlist  = {'spoke':[], 'medbeta':[], 'highbeta':[]}
    cavnum = 0
    Pblist   = {'spoke':[], 'medbeta':[], 'highbeta':[]}
    Preflist = {'spoke':[], 'medbeta':[], 'highbeta':[]}
    Pglist   = {'spoke':[], 'medbeta':[], 'highbeta':[]}
    
    #Qlopt = {'spoke': 2.37e5, 'medbeta': 8e5, 'highbeta':7.5e5}
    
    spokedata    = csv.reader(open('spokenewshort.csv', 'rU'), dialect=csv.excel_tab, delimiter=',')
    medbetadata  = csv.reader(open('lowbetanew.csv', 'rU'), dialect=csv.excel_tab, delimiter=',')
    highbetadata = csv.reader(open('highbetanew.csv', 'rU'), dialect=csv.excel_tab, delimiter=',')
    for row in spokedata:
        A = cavity(f=352.21e6, delF=nan, I_b0=50e-3, I_b=50e-3, V_cav=float(row[0])*float(row[3]), beamph=float(row[1]), RoverQ=500.0, 
                Ql=Qlopt['spoke'], beta_0=0.46, beta=float(row[2]), Eampfile='Ezvszspoke.dat')
        #fig1111 = plt.figure(cavnum)
        #A.drawvecs(fig1111)
        #f = "pubspokevecs_%d.pdf" % cavnum
        #fig1111.savefig(f)
        print "Cav volts = %e, R/Q = %f, R/Q scaling = %f" % (abs(A.V_cav), A.RoverQ, A.RQscaling)
        #print "cavity(f=%e, delF=nan, I_b0=50e-3, I_b=50e-3, V_cav=%e, beamph=%f, RoverQ=%d, Ql=%e, beta_0=%f, beta=%f)" % (
        #        A.f, abs(A.V_cav), A.beamph, A.RoverQ, A.Ql, A.beta_0, A.beta)
        cavnum += 1
        cavlist['spoke'].append(cavnum)
        Preflist['spoke'].append( A.Pref )
        Pblist['spoke'].append( A.Pbeam )
        Pglist['spoke'].append( A.Pg )
    
    for row in medbetadata:
        A = cavity(f=704.42e6, delF=nan, I_b0=50e-3, I_b=50e-3, V_cav=float(row[0])*float(row[3]), beamph=float(row[1]), RoverQ=300.0, 
                Ql=Qlopt['medbeta'], beta_0=0.70, beta=float(row[2]))
        #fig1111 = plt.figure(cavnum)
        #A.drawvecs(fig1111)
        #f = "pubmedbetavecs_%d.pdf" % cavnum
        #fig1111.savefig(f)
        print "Cav volts = %e, R/Q = %f, R/Q scaling = %f" % (abs(A.V_cav), A.RoverQ, A.RQscaling)
        cavnum += 1
        cavlist['medbeta'].append(cavnum)
        Preflist['medbeta'].append( A.Pref )
        Pblist['medbeta'].append( A.Pbeam )
        Pglist['medbeta'].append( A.Pg )
    
    for row in highbetadata:
        A = cavity(f=704.42e6, delF=nan, I_b0=50e-3, I_b=50e-3, V_cav=float(row[0])*float(row[3]), beamph=float(row[1]), RoverQ=477.0, 
                Ql=Qlopt['highbeta'], beta_0=0.92, beta=float(row[2]))
        #fig1111 = plt.figure(cavnum)
        #A.drawvecs(fig1111)
        #f = "pubhighbetavecs_%d.pdf" % cavnum
        #fig1111.savefig(f)
        #print "Cav volts = %e, R/Q = %f, R/Q scaling = %f" % (abs(A.V_cav), A.RoverQ, A.RQscaling)
        cavnum += 1
        cavlist['highbeta'].append(cavnum)
        Preflist['highbeta'].append( A.Pref )
        Pblist['highbeta'].append( A.Pbeam )
        Pglist['highbeta'].append( A.Pg )
    
    W2kW = lambda x: x/1e3
    makeneg = lambda x: -x
    fig = plt.figure(1)
    ax = fig.add_subplot(211)
    ax.plot(cavlist['spoke'], map(W2kW, Pglist['spoke']), '-', label='Spoke forward')
    ax.plot(cavlist['spoke'], map(W2kW, Pblist['spoke']), 'x', label='Spoke to beam')
    ax.plot(cavlist['medbeta'], map(W2kW, Pglist['medbeta']), '-', label='Med Beta forward')
    ax.plot(cavlist['medbeta'], map(W2kW, Pblist['medbeta']), 'x', label='Med Beta to beam')
    ax.plot(cavlist['highbeta'], map(W2kW, Pglist['highbeta']), '-', label='High Beta forward')
    ax.plot(cavlist['highbeta'], map(W2kW, Pblist['highbeta']), 'x', label='High Beta to beam')
    ax.set_xlim(right=cavlist['highbeta'][-1]+1)
    ax.legend(loc=0)
    ax.set_xlabel('Resonator number')
    ax.set_ylabel('Power (kW)')
    ax.grid()
    ax.set_ylim(bottom=0)
    
    ax2 = fig.add_subplot(234)
    ax3 = fig.add_subplot(235)
    ax4 = fig.add_subplot(236)
    ax2.plot(cavlist['spoke'], map(W2kW, Preflist['spoke']), '-x', label="""Spoke reflected""")
    ax2.set_title('Ql = %0.3e' % Qlopt['spoke'])
    ax3.set_title('Ql = %0.3e' % Qlopt['medbeta'])
    ax4.set_title('Ql = %0.3e' % Qlopt['highbeta'])
    ax3.plot(cavlist['medbeta'], map(W2kW, Preflist['medbeta']), '-x', label='Med Beta reflected')
    ax4.plot(cavlist['highbeta'], map(W2kW, Preflist['highbeta']), '-x', label='High Beta reflected')
    for axis in (ax2, ax3, ax4):
        axis.legend(loc=0)
        axis.set_xlabel('Resonator number')
        axis.set_ylabel('Power (kW)')
        axis.grid()
        axis.set_ylim(bottom=0)
    
    nompower = {'spoke':0, 'medbeta':0, 'highbeta':0}
    optpower = {'spoke':0, 'medbeta':0, 'highbeta':0}
    
    optpower['spoke']    = linacreflected(Qlopt['spoke'], 'spokenewshort.csv', 352.21e6, nan, 50e-3, 50e-3, 500.0, 0.46, 'Ezvszspoke.dat', False)
    optpower['medbeta']  = linacreflected(Qlopt['medbeta'], 'lowbetanew.csv',  704.42e6, nan, 50e-3, 50e-3, 300.0, 0.70, nan, False)
    optpower['highbeta'] = linacreflected(Qlopt['highbeta'], 'highbetanew.csv', 704.42e6, nan, 50e-3, 50e-3, 477.0, 0.92, nan, False)
    
    Qlopt = {'spoke': 2.37e5, 'medbeta': 8e5, 'highbeta':7.5e5}
    nompower['spoke']    = linacreflected(Qlopt['spoke'], 'spokenewshort.csv', 352.21e6, nan, 50e-3, 50e-3, 500.0, 0.46, 'Ezvszspoke.dat', False)
    nompower['medbeta']  = linacreflected(Qlopt['medbeta'], 'lowbetanew.csv',  704.42e6, nan, 50e-3, 50e-3, 300.0, 0.70, nan, False)
    nompower['highbeta'] = linacreflected(Qlopt['highbeta'], 'highbetanew.csv', 704.42e6, nan, 50e-3, 50e-3, 477.0, 0.92, nan, False)
    
    totalpower = {'nominal':nompower['spoke']+nompower['medbeta']+nompower['highbeta'],
            'optimised':optpower['spoke']+optpower['medbeta']+optpower['highbeta']}
    
    print "Nominal:   Total reflected power = %0.1f kW" % (totalpower['nominal']/1e3)
    print "Optimised: Total reflected power = %0.1f kW" % (totalpower['optimised']/1e3)
    print "Difference: %0.1f kW" % ((totalpower['nominal'] - totalpower['optimised'])/1e3)
    #print "Total forward power: %f kW" % ((sum(Pforlist['spoke'])+sum(Pforlist['medbeta'])+sum(Pforlist['highbeta']))/1e3)
    
    plt.show()
    
